#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#include "hyper/Hyper.h"

#include <iostream>

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;
using namespace vector;
using namespace mlir::arith;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

class HyperAllocLowering : public OpRewritePattern<hyper::AllocOp> {
public:
  using OpRewritePattern<hyper::AllocOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(hyper::AllocOp op,
                                PatternRewriter &rewriter) const override {
    StringRef dev = op.getDevice();
    auto memrefType = op.getType().cast<MemRefType>();
    auto dynamicSizes = op.getDynamicSizes();
    auto symbolOperands = op.getSymbolOperands();

    if (dev.equals("cpu") || dev.equals("cgra")) {
      if (memrefType.getNumDynamicDims() == 0) {
        rewriter.replaceOpWithNewOp<memref::AllocOp>(op,
                                                     /*memrefType=*/memrefType);
      } else {
        rewriter.replaceOpWithNewOp<memref::AllocOp>(
            op, /*memrefType=*/memrefType, /*dynamicSizes=*/dynamicSizes,
            /*symbolOperands=*/symbolOperands);
      }
    } else if (dev.equals("gpu")) {
      if (memrefType.getNumDynamicDims() == 0) {
        rewriter.replaceOpWithNewOp<gpu::AllocOp>(
            op, /*memref=*/memrefType, /*asyncToken=*/nullptr,
            /*asyncDependencies=*/ValueRange{}, /*dynamicSizes=*/ValueRange{},
            /*symbolOperands=*/symbolOperands, /*hostShared=*/false);
      } else {
        rewriter.replaceOpWithNewOp<gpu::AllocOp>(
            op, /*memref=*/memrefType, /*asyncToken=*/nullptr,
            /*asyncDependencies=*/ValueRange{}, /*dynamicSizes=*/dynamicSizes,
            /*symbolOperands=*/symbolOperands, /*hostShared=*/false);
      }
    } else {
      return rewriter.notifyMatchFailure(
          op, "unsupported device for hyper.alloc now.");
    }

    return success();
  }
};

class HyperDeallocLowering : public OpRewritePattern<hyper::DeallocOp> {
public:
  using OpRewritePattern<hyper::DeallocOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(hyper::DeallocOp op,
                                PatternRewriter &rewriter) const override {
    StringRef dev = op.getDevice();
    auto memref = op.getMemref();

    if (dev.equals("cpu")) {
      rewriter.replaceOpWithNewOp<memref::DeallocOp>(op, /*memref=*/memref);
    } else if (dev.equals("gpu")) {
      rewriter.replaceOpWithNewOp<gpu::DeallocOp>(
          op, /*asyncToken=*/std::nullopt, /*asyncDependencies=*/ValueRange{},
          /*memref=*/memref);
    } else {
      return rewriter.notifyMatchFailure(
          op, "unsupported device for hyper.dealloc now.");
    }

    return success();
  }
};

} // namespace

void populateLowerHyperPatterns(RewritePatternSet &patterns) {
  patterns.add<HyperAllocLowering>(patterns.getContext());
  patterns.add<HyperDeallocLowering>(patterns.getContext());
}

namespace {
class LowerHyperPass
    : public PassWrapper<LowerHyperPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerHyperPass)
  LowerHyperPass() = default;
  LowerHyperPass(const LowerHyperPass &) {}

  StringRef getArgument() const final { return "lower-hyper"; }
  StringRef getDescription() const final {
    return "Lower hyper ops to device-specific ops.";
  }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect, mlir::LLVM::LLVMDialect,
                    memref::MemRefDialect, arith::ArithDialect,
                    vector::VectorDialect, gpu::GPUDialect, scf::SCFDialect,
                    hyper::HyperDialect>();
  }
};
} // end anonymous namespace.

void LowerHyperPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<func::FuncDialect, mlir::LLVM::LLVMDialect,
                         memref::MemRefDialect, arith::ArithDialect,
                         gpu::GPUDialect, vector::VectorDialect,
                         scf::SCFDialect>();
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();

  RewritePatternSet patterns(context);
  populateLowerHyperPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerHyperPass() { PassRegistration<LowerHyperPass>(); }
